import torch
import numpy as np
from torch import nn 
from torch.autograd import Variable
import torch.nn.functional as F
from torch.utils.data import DataLoader
from datetime import datetime

def cov_3x3(in_channel, out_channel, stride = 1):
    return nn.Conv2d(in_channel, out_channel, 3, stride=stride, padding=1, bias=False)

class res_block(nn.Module):
    def __init__(self, in_channel, out_channel, same_shape = True):
        super(res_block,self).__init__()
        self.same_shape = same_shape
        stride = 1 if same_shape else 2

        self.conv1 = cov_3x3(in_channel, out_channel, stride=stride)
        self.bn1 = nn.BatchNorm2d(out_channel)
        self.conv2 = cov_3x3(out_channel, out_channel)
        self.bn2 = nn.BatchNorm2d(out_channel)
        if not same_shape:
            self.conv3 = nn.Conv2d(in_channel, out_channel, 1, stride=stride)

    def forward(self,x):
        out = self.conv1(x)
        out = F.relu(self.bn1(out))
        out = self.conv2(out)
        out = F.relu(self.bn2(out))
        if not self.same_shape:
            x = self.conv3(x)
        return F.relu(x + out)
    
class net_sensor(nn.Module):
    def __init__(self, in_channel, verbose = False):
        super(net_sensor, self).__init__()
        
        self.verbose = verbose
        #版本1
        self.block1 = cov_3x3(in_channel, 16)
        self.block2 = res_block(16, 16)
        self.block3 = res_block(16, 16)
        self.block4 = cov_3x3(16, 1)
        self.block5 = nn.AvgPool2d(32,1)



    def forward(self,x):
        x = self.block1(x)
        if self.verbose:
            print('block1 output {}'.format(x.shape))
        x = F.relu(x)
        x = self.block2(x)
        if self.verbose:
            print('block2 output {}'.format(x.shape))        
        x = self.block3(x)
        if self.verbose:
            print('block3 output {}'.format(x.shape))   
        x = self.block4(x)
        if self.verbose:
            print('block4 output {}'.format(x.shape))   
        x = F.relu(x)
        x = self.block5(x)
        if self.verbose:
            print('block5 output {}'.format(x.shape))           
        return x


class image_rec(nn.Module):
    def __init__(self, in_channel, verbose = False):
        super(image_rec, self).__init__()
        self.verbose = verbose
        #版本1，深度不够
        # self.block1 = cov_3x3(in_channel, 64)
        # self.block2 = nn.MaxPool2d(2, 2)
        # self.block3 = res_block(64, 64)
        # self.block4 = nn.Dropout(0.2)
        # self.block5 = res_block(64, 128, False)
        # self.block6 = nn.AvgPool2d(8, 1)
        # self.block7 = nn.Flatten(1)
        # self.classifier = nn.Linear(128, 9)

        self.block1 = nn.Conv2d(in_channel, 64, 7, 1, 3)
        self.block2 = nn.MaxPool2d(2,2)
        self.block3 = nn.Sequential(
            res_block(64, 64),
            res_block(64, 128, False)
        )
        self.dropout = nn.Dropout(0.2)
        self.block4 = nn.Sequential(
            res_block(128, 128),
            res_block(128, 256, False)
        )
        self.block5 = nn.Sequential(
            res_block(256, 256),
            res_block(256, 512, False)
        )
        self.block6 = nn.AvgPool2d(2,1)
        self.block7 = nn.Flatten(1)
        self.classifier = nn.Linear(512,9)

    def forward(self, x):
        x = self.block1(x)
        if self.verbose:
            print("block1 output {}".format(x.shape))
        x = self.block2(x)
        if self.verbose:
            print("block2 output {}".format(x.shape)) 
        x = self.block3(x)
        if self.verbose:
            print("block3 output {}".format(x.shape))   
        x = self.block4(x)
        if self.verbose:
            print("block4 output {}".format(x.shape))
        x = self.block5(x)
        if self.verbose:
            print("block5 output {}".format(x.shape))
        x = self.block6(x)
        if self.verbose:
            print("block6 output {}".format(x.shape))
        x = self.block7(x)
        if self.verbose:
            print("block7 output {}".format(x.shape))            
        x = self.classifier(x) 
        if self.verbose:
            print("classifier output {}".format(x.shape))   
        return x          
       
def tf_resnet(x):
    x = x.reshape(-1,1,32,32)
    x = np.array(x, dtype = 'float32')/4
    x = (x-0.5)/0.5
    x = torch.from_numpy(x)
    return x

def tf_label(x):
    x = x.reshape(-1,1,1,1)
    x = np.array(x, dtype = 'float32')/1000
    x = torch.from_numpy(x)
    return x   

def get_acc(output, label):
    total = output.shape[0]
    _, pred_label = output.max(1)
    num_correct = (pred_label == label).sum().data
    return num_correct / total

def train_resnet(net, train_data, test_data , num_epochs, optimizer, criterion):
    if torch.cuda.is_available():
        net = net.cuda()
    prev_time = datetime.now()
    for epoch in range(num_epochs):
        train_loss = 0
        train_acc = 0
        net = net.train()
        for data1 in train_data:
            if torch.cuda.is_available():
                data_train = Variable(tf_resnet(data1[:,0:-1]).cuda(), volatile = True)
                label_train = Variable(data1[:,-1].type(torch.int64).cuda(), volatile = True)
            else:
                data_train = Variable(tf_resnet(data1[:,0:-1]), volatile = True)
                label_train = Variable(data1[:,-1].type(torch.int64), volatile = True)
            #forward
            output = net(data_train)
            loss = criterion(output, label_train)
            #backward
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()
            train_loss += loss
            train_acc += get_acc(output, label_train)
        cur_time = datetime.now()
        h, remainder = divmod((cur_time - prev_time).seconds, 3600)
        m, s = divmod(remainder, 60)
        time_str = "Time %02d:%02d:%02d" % (h, m, s)
        if test_data is not None:
            test_loss = 0
            test_acc = 0
            net = net.eval()
            for data1 in test_data:
                if torch.cuda.is_available():
                    data_train = Variable(tf_resnet(data1[:,0:-1]).cuda(), volatile = True)
                    label_train = Variable(data1[:,-1].type(torch.int64).cuda(), volatile = True)
                else:
                    data_train = Variable(tf_resnet(data1[:,0:-1]), volatile = True)
                    label_train = Variable(data1[:,-1].type(torch.int64), volatile = True)
                #forward
                output = net(data_train)
                loss = criterion(output, label_train)
                test_loss += loss
                test_acc += get_acc(output, label_train)
            epoch_str = (
                "epoch:%d. Train Loss:%f, Train acc:%f, Test Loss:%f, Test acc:%f"%(epoch, train_loss/len(train_data), train_acc/len(train_data), test_loss/len(test_data), test_acc/len(test_data))
            )
        else:
            epoch_str = (
                "epoch:%d. Train Loss:%f, Train acc:%f,"%(epoch, train_loss/len(train_data), train_acc/len(train_data))
            )
        prev_time = cur_time
        print(time_str + epoch_str)


# net1 = image_rec(1, True)
# test_x = Variable(torch.zeros(1,1,32,32))
# test_y = net1(test_x)
    

    




